# encoding: utf-8
import json
import logging
import traceback
import numpy as np
from typing import List
from modelscope.pipelines import pipeline

LOGGER = logging.getLogger("命名实体识别")


# Extend the JSONEncoder class
class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)


class NamedEntityExtractor(object):

    def __init__(self, pretrained_model):
        self.pipeline = pipeline('named-entity-recognition', pretrained_model)

    def extract(self, text: str):
        try:
            result = self.pipeline(text)
            LOGGER.info(result)
            # NP序列化问题简易解决办法
            result = json.loads(json.dumps(result, cls=NpEncoder))
            return result
        except Exception as e:
            LOGGER.error(traceback.format_exc())

        return []

    def extract_batch(self, texts: List[str]):
        try:
            result = self.pipeline(texts, batch_size=128)
            LOGGER.info(result)
            # NP序列化问题简易解决办法
            result = json.loads(json.dumps(result, cls=NpEncoder))
            return result
        except Exception as e:
            LOGGER.error(traceback.format_exc())

        return []
